####!/usr/bin/env python

"""
Functions for data IO for neural network training.
"""

from __future__ import print_function
import argparse
import sys
import os
import time
import matplotlib.pyplot as plt
from operator import add
import math
import numpy as np
import random
import imblearn as ib
import seaborn as sns
import sklearn
from imblearn.over_sampling import RandomOverSampler
from imblearn.over_sampling import SMOTE
from keras.layers import Input, Dense, Conv1D, Dropout


def mkdir(outdir):
    if not os.path.exists(outdir):
        os.makedirs(outdir)


def clean_file(file_path):
    if os.path.isfile(file_path):
        os.remove(file_path)


def enc_list_bl_max_len(aa_seqs, blosum, max_seq_len):
    '''
    blosum encoding of a list of amino acid sequences with padding 
    to a max length

    parameters:
        - aa_seqs : list with AA sequences
        - blosum : dictionnary: key= AA, value= blosum encoding
        - max_seq_len: common length for padding
    returns:
        - enc_aa_seq : list of np.ndarrays containing padded, encoded amino acid sequences
    '''

    # encode sequences:
    sequences = []
    for seq in aa_seqs:
        e_seq = np.zeros((len(seq), len(blosum["A"])))
        count = 0
        for aa in seq:
            if aa in blosum:
                e_seq[count] = blosum[aa]
                count += 1
            else:
                sys.stderr.write("Unknown amino acid in peptides: " + aa + ", encoding aborted!\n")
                sys.exit(2)

        sequences.append(e_seq)

    # pad sequences:
    # max_seq_len = max([len(x) for x in aa_seqs])
    n_seqs = len(aa_seqs)
    n_features = sequences[0].shape[1]

    enc_aa_seq = np.zeros((n_seqs, max_seq_len, n_features))
    for i in range(0, n_seqs):
        enc_aa_seq[i, :sequences[i].shape[0], :n_features] = sequences[i]

    return enc_aa_seq


blosum50_20aa = {
    'A': np.array((5, -2, -1, -2, -1, -1, -1, 0, -2, -1, -2, -1, -1, -3, -1, 1, 0, -3, -2, 0)),
    'R': np.array((-2, 7, -1, -2, -4, 1, 0, -3, 0, -4, -3, 3, -2, -3, -3, -1, -1, -3, -1, -3)),
    'N': np.array((-1, -1, 7, 2, -2, 0, 0, 0, 1, -3, -4, 0, -2, -4, -2, 1, 0, -4, -2, -3)),
    'D': np.array((-2, -2, 2, 8, -4, 0, 2, -1, -1, -4, -4, -1, -4, -5, -1, 0, -1, -5, -3, -4)),
    'C': np.array((-1, -4, -2, -4, 13, -3, -3, -3, -3, -2, -2, -3, -2, -2, -4, -1, -1, -5, -3, -1)),
    'Q': np.array((-1, 1, 0, 0, -3, 7, 2, -2, 1, -3, -2, 2, 0, -4, -1, 0, -1, -1, -1, -3)),
    'E': np.array((-1, 0, 0, 2, -3, 2, 6, -3, 0, -4, -3, 1, -2, -3, -1, -1, -1, -3, -2, -3)),
    'G': np.array((0, -3, 0, -1, -3, -2, -3, 8, -2, -4, -4, -2, -3, -4, -2, 0, -2, -3, -3, -4)),
    'H': np.array((-2, 0, 1, -1, -3, 1, 0, -2, 10, -4, -3, 0, -1, -1, -2, -1, -2, -3, 2, -4)),
    'I': np.array((-1, -4, -3, -4, -2, -3, -4, -4, -4, 5, 2, -3, 2, 0, -3, -3, -1, -3, -1, 4)),
    'L': np.array((-2, -3, -4, -4, -2, -2, -3, -4, -3, 2, 5, -3, 3, 1, -4, -3, -1, -2, -1, 1)),
    'K': np.array((-1, 3, 0, -1, -3, 2, 1, -2, 0, -3, -3, 6, -2, -4, -1, 0, -1, -3, -2, -3)),
    'M': np.array((-1, -2, -2, -4, -2, 0, -2, -3, -1, 2, 3, -2, 7, 0, -3, -2, -1, -1, 0, 1)),
    'F': np.array((-3, -3, -4, -5, -2, -4, -3, -4, -1, 0, 1, -4, 0, 8, -4, -3, -2, 1, 4, -1)),
    'P': np.array((-1, -3, -2, -1, -4, -1, -1, -2, -2, -3, -4, -1, -3, -4, 10, -1, -1, -4, -3, -3)),
    'S': np.array((1, -1, 1, 0, -1, 0, -1, 0, -1, -3, -3, 0, -2, -3, -1, 5, 2, -4, -2, -2)),
    'T': np.array((0, -1, 0, -1, -1, -1, -1, -2, -2, -1, -1, -1, -1, -2, -1, 2, 5, -3, -2, 0)),
    'W': np.array((-3, -3, -4, -5, -5, -1, -3, -3, -3, -3, -2, -3, -1, 1, -4, -4, -3, 15, 2, -3)),
    'Y': np.array((-2, -1, -2, -3, -3, -1, -2, -3, 2, -1, -1, -2, 0, 4, -3, -2, -2, 2, 8, -1)),
    'V': np.array((0, -3, -3, -4, -1, -3, -3, -4, -4, 4, 1, -3, 1, -1, -3, -2, 0, -3, -1, 5))
}


# 划分训练集、验证集；通过过采样将训练集分成批
def spilt_data(train_input, train_y, EPOCH, BATCH_SIZE, val_rate):
    # 一个输入列表，每个epoch从中取值一次，得到一个batch_train_input
    ep_train_input = []
    ep_train_y = []
    # 区分正负样本
    pos_index = np.nonzero(train_y)[0]
    neg_index = np.nonzero(~train_y)[0]
    # 划分验证集
    valid_pos_index_index = np.random.randint(0, np.size(pos_index) - 1, int(BATCH_SIZE * val_rate / 2))
    valid_neg_index_index = np.random.randint(0, np.size(neg_index) - 1, int(BATCH_SIZE * val_rate / 2))
    valid_pos_index = pos_index[valid_pos_index_index]
    valid_neg_index = neg_index[valid_neg_index_index]
    valid_index_batches = valid_pos_index.tolist() + valid_neg_index.tolist()
    valid_batches = [train_input[0][np.array(valid_index_batches)], train_input[1][np.array(valid_index_batches)],
                     train_input[2][np.array(valid_index_batches)]]
    valid_y_batches = train_y[valid_index_batches]
    ep_val_data = (valid_batches, valid_y_batches)
    # 划分训练集
    pos_index = np.delete(pos_index, valid_pos_index_index)
    neg_index = np.delete(neg_index, valid_neg_index_index)

    for j in range(int(len(train_y) / BATCH_SIZE)):
        # 过采样得到正反例相等的batch
        pos_index_batches = random.sample(pos_index.tolist(), int(0.5 * BATCH_SIZE))
        neg_index_batches = random.sample(neg_index.tolist(), BATCH_SIZE - int(0.5 * BATCH_SIZE))
        index_bathes = pos_index_batches + neg_index_batches
        train_batches = [train_input[0][np.array(index_bathes)], train_input[1][np.array(index_bathes)],
                         train_input[2][np.array(index_bathes)]]
        y_train_batches = train_y[np.array(index_bathes)]
        ep_train_input.append(train_batches)
        ep_train_y.append(y_train_batches)

    return ep_train_input, ep_train_y, ep_val_data


# TODO 过采样
def smote_over_sample(x_train, y_train):
    sos = SMOTE(sampling_strategy=0.1, random_state=10, k_neighbors=5, n_jobs=-1)
    x_sos, y_sos = sos.fit_resample(x_train, y_train)
    return x_sos, y_sos


# TODO 绘图
# 输入指标名和历史对象，可以得到模型各期的指标得分
def get_metric_from_history(history, metric):
    metric = str(metric)
    res = []
    for epoch_his in history:
        ep_res = []
        try:
            for batch_his in epoch_his:
                ep_res.append(batch_his[metric])
            res.append(np.average(ep_res))
        except KeyError as e:
            print("!!指标" + metric + "不存在!!")
            break

    return res


# 根据bound来计算四个指标。bound一般是0.5
def cal_metric(y, score_y, bound):
    pred_y = np.copy(score_y)
    pred_y[pred_y > bound] = 1
    pred_y[pred_y < bound] = 0
    acc = sklearn.metrics.accuracy_score(y, pred_y)
    recall = sklearn.metrics.recall_score(y, pred_y)
    precision = sklearn.metrics.precision_score(y, pred_y)
    roc_auc = sklearn.metrics.roc_auc_score(y, score_y)
    return acc, recall, precision, roc_auc


def plot_test_graphs(index_array, index_name, pic_output_path):
    plt.plot(range(len(index_array)), index_array, 'b', label=index_name)
    plt.title('Testing_' + str(index_name))
    plt.xlabel('Epochs')
    plt.ylabel(str(index_name))
    plt.legend()
    plt.show()
    pic_path = pic_output_path + str(index_name) + '.png'
    clean_file(pic_path)
    plt.savefig(pic_path)


def plot_graphs(list1, list2, list_name, pic_path):
    plt.plot(list1, 'r')
    plt.plot(list2, 'b')
    plt.xlabel("Epoch")
    plt.ylabel(list_name)
    plt.legend([list_name, 'val_' + list_name])
    plt.show()
    pic_path = pic_path + list_name + '__' + 'val_' + list_name
    clean_file(pic_path)
    plt.savefig(pic_path)
